package com.yang.ChatOllama.controller;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Media;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.core.io.ClassPathResource;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.List;

@RestController
public class OllamaController {
    @Resource
    private OllamaChatClient ollamaChatClient;


    @RequestMapping("/ai/ollama")
    public Object chat(@RequestParam(value = "msg") String msg) {
        ChatResponse response = ollamaChatClient.call(
                new Prompt(msg,
                        OllamaOptions.create()
                                .withModel("qwen:14b")
                                .withTemperature(0.8f)
                )
        );
        System.out.print(response.getResult().getOutput());
        return response.getResult().getOutput();
    }
    @GetMapping("/ai/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        Prompt prompt = new Prompt(new UserMessage(message),
                OllamaOptions.create().withModel("qwen:14b"));
        return ollamaChatClient.stream(prompt);
    }

    /**
     * 本次使用llava:7b 硬件有限，怕跑的很慢 所以不返回结果，直接打印
     */
    @GetMapping("/ai/test")
    public void test() throws IOException {
        byte[] imageData = new ClassPathResource("/multimodal.test.png").getContentAsByteArray();

        var userMessage = new UserMessage("描述一下图片内容?",
                List.of(new Media(MimeTypeUtils.IMAGE_PNG, imageData)));

        ChatResponse response = ollamaChatClient.call(
                new Prompt(List.of(userMessage),
                        OllamaOptions.create().withModel("llava:7b").withMainGPU(1)
                )
        );
        System.out.println(response.getResult().getOutput().getContent());
    }

    //更多的图片相关的接口开发可以参考MultiModel项目的内容
}
